import pandas as pd
import os
from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED, as_completed

EXCLUDE_FILES = ['通行设施.csv', '室内设施.csv', '地名地址信息.csv', '事件活动.csv']


executor = ThreadPoolExecutor(max_workers=2)
types = pd.read_csv('/home/fredyvia/dbf_data_solve/res_types.csv')
types_tmp = types.groupby(
    ["SUBTYPE", "CATEGORY"]).size().to_frame("COUNT")
types_tmp = types_tmp[types_tmp["COUNT"] != 1]
assert types_tmp.empty


def get_all_csv(source_path, dest_path):
    res = []
    for fileName in os.listdir(source_path):
        source_file_path = source_path+os.sep+fileName
        if fileName.endswith(".csv") and fileName not in EXCLUDE_FILES:
            res.append(
                (source_file_path, dest_path+os.sep +
                 os.path.split(source_file_path)[-1]
                 )
            )
        elif os.path.isdir(source_file_path):
            res.extend(get_all_csv(source_path+os.sep +
                       fileName, dest_path+os.sep+fileName))
    return res


def solve(source_file, dest_file):
    print("solving", source_file)
    dest_dir = os.path.dirname(dest_file)
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)
    # tag = os.path.splitext(os.path.split(source_file)[-1])[0]
    # types_tmp = types[types["TAG"] == tag].drop(columns=["TAG"])
    # if(types_tmp.empty):
    #     unsolved_files.append(source_file)
    #     print("unsolved file:", source_file)
    #     continue
    df_tmp = pd.read_csv(source_file)
    df_tmp = df_tmp.merge(types, on=["SUBTYPE", "CATEGORY"], how="left")
    df_tmp.to_csv("test.csv")
    df_tmp["大类"] = df_tmp["中类"].map(lambda tt: tt[0])
    df_tmp.to_csv(dest_file, index=False)
    return source_file


# print(types)
files = get_all_csv("csv/2017", "csv_type2/2017")
# print(files)
# files = [("csv/2012/地名地址信息.csv", "csv_type/2012/地名地址信息.csv")]
# unsolved_files = []
all_task = [executor.submit(lambda p:solve(*p),[source, dest])
                    for source, dest in files]
# wait(all_task, return_when=ALL_COMPLETED)
for future in as_completed(all_task):
    data = future.result()
    print("{} success".format(data))
executor.shutdown()
